#include <iostream>
#include <cmath>
#include <vector>
#include <cstdlib>
#include <ctime>

using namespace std;

const int NUM_ANTS = 10;
const int NUM_ITERATIONS = 100;
const double ALPHA = 1.0; // 蚂蚁信息素重要程度
const double BETA = 2.0; // 启发信息重要程度
const double RHO = 0.5; // 信息素挥发速度
const double Q = 100.0; // 信息素增加强度
const double INITIAL_PHEROMONE = 1.0;

class AntColony {
private:
    int numCities;
    vector<vector<double>> distances;
    vector<vector<double>> pheromones;
    vector<int> bestTour;
    double bestTourLength;

public:
    AntColony(int numCities, vector<vector<double>>& distances) : numCities(numCities), distances(distances) {
        pheromones.resize(numCities, vector<double>(numCities, INITIAL_PHEROMONE));
    }

    void run() {
        srand(time(nullptr));
        for (int iter = 0; iter < NUM_ITERATIONS; ++iter) {
            vector<vector<int>> antTours(NUM_ANTS);
            vector<double> tourLengths(NUM_ANTS);

            for (int ant = 0; ant < NUM_ANTS; ++ant) {
                vector<int> tour;
                vector<bool> visited(numCities, false);
                int startCity = rand() % numCities;
                tour.push_back(startCity);
                visited[startCity] = true;
                for (int i = 0; i < numCities - 1; ++i) {
                    int city = selectNextCity(tour.back(), visited);
                    tour.push_back(city);
                    visited[city] = true;
                }
                antTours[ant] = tour;
                tourLengths[ant] = calculateTourLength(tour);
            }

            updatePheromones(antTours, tourLengths);

            if (iter == 0 || tourLengths[bestAnt(tourLengths)] < bestTourLength) {
                bestTourLength = tourLengths[bestAnt(tourLengths)];
                bestTour = antTours[bestAnt(tourLengths)];
            }
        }

        cout << "Best tour length: " << bestTourLength << endl;
        cout << "Best tour: ";
        for (int city : bestTour) {
            cout << city << " ";
        }
        cout << endl;
    }

private:
    double calculateTourLength(const vector<int>& tour) {
        double length = 0;
        for (int i = 0; i < numCities - 1; ++i) {
            length += distances[tour[i]][tour[i + 1]];
        }
        length += distances[tour[numCities - 1]][tour[0]]; // 回到起点
        return length;
    }

    int selectNextCity(int currentCity, const vector<bool>& visited) {
        double totalProbability = 0.0;
        vector<double> probabilities(numCities, 0.0);
        for (int i = 0; i < numCities; ++i) {
            if (!visited[i]) {
                probabilities[i] = pow(pheromones[currentCity][i], ALPHA) * pow(1.0 / distances[currentCity][i], BETA);
                totalProbability += probabilities[i];
            }
        }
        double random = static_cast<double>(rand()) / RAND_MAX * totalProbability;
        double sum = 0.0;
        for (int i = 0; i < numCities; ++i) {
            if (!visited[i]) {
                sum += probabilities[i];
                if (sum >= random) {
                    return i;
                }
            }
        }
        return -1; // Should not reach here
    }

    void updatePheromones(const vector<vector<int>>& antTours, const vector<double>& tourLengths) {
        for (int i = 0; i < numCities; ++i) {
            for (int j = 0; j < numCities; ++j) {
                pheromones[i][j] *= (1.0 - RHO);
            }
        }
        for (int ant = 0; ant < NUM_ANTS; ++ant) {
            double contribution = Q / tourLengths[ant];
            for (int i = 0; i < numCities - 1; ++i) {
                pheromones[antTours[ant][i]][antTours[ant][i + 1]] += contribution;
            }
            pheromones[antTours[ant][numCities - 1]][antTours[ant][0]] += contribution; // 回到起点
        }
    }

    int bestAnt(const vector<double>& tourLengths) {
        double minTourLength = tourLengths[0];
        int bestAntIndex = 0;
        for (int i = 1; i < NUM_ANTS; ++i) {
            if (tourLengths[i] < minTourLength) {
                minTourLength = tourLengths[i];
                bestAntIndex = i;
                if (minTourLength == 0) {
                }
            }
        }
        return bestAntIndex;
    }
};

int main() {
    vector<vector<double>> distances = {
        {0, 10, 15, 20},
        {10, 0, 35, 25},
        {15, 35, 0, 30},
        {20, 25, 30, 0}
    };
    int numCities = distances.size();
    AntColony antColony(numCities, distances);
    antColony.run();
    return 0;
}